import os
import logging
from typing import Dict, Any, Optional
from dataclasses import dataclass
from contextlib import contextmanager
import pymysql
from pymysql.cursors import DictCursor
from pymysql import pooling


@dataclass
class DatabaseConfig:
    """数据库配置类"""
    host: str = 'localhost'
    port: int = 3306
    user: str = 'root'
    password: str = ''
    database: str = 'eeg_signal_db'
    charset: str = 'utf8mb4'
    max_connections: int = 20
    min_connections: int = 5
    connection_timeout: int = 30
    read_timeout: int = 30
    write_timeout: int = 30
    
    @classmethod
    def from_env(cls) -> 'DatabaseConfig':
        """从环境变量创建配置"""
        return cls(
            host=os.getenv('MYSQL_HOST', 'localhost'),
            port=int(os.getenv('MYSQL_PORT', '3306')),
            user=os.getenv('MYSQL_USER', 'root'),
            password=os.getenv('MYSQL_PASSWORD', ''),
            database=os.getenv('MYSQL_DATABASE', 'eeg_signal_db'),
            charset=os.getenv('MYSQL_CHARSET', 'utf8mb4'),
            max_connections=int(os.getenv('MYSQL_MAX_CONNECTIONS', '20')),
            min_connections=int(os.getenv('MYSQL_MIN_CONNECTIONS', '5')),
            connection_timeout=int(os.getenv('MYSQL_CONNECTION_TIMEOUT', '30')),
            read_timeout=int(os.getenv('MYSQL_READ_TIMEOUT', '30')),
            write_timeout=int(os.getenv('MYSQL_WRITE_TIMEOUT', '30'))
        )
    
    def get_connection_params(self) -> Dict[str, Any]:
        """获取连接参数"""
        return {
            'host': self.host,
            'port': self.port,
            'user': self.user,
            'password': self.password,
            'database': self.database,
            'charset': self.charset,
            'cursorclass': DictCursor,
            'autocommit': True,
            'connect_timeout': self.connection_timeout,
            'read_timeout': self.read_timeout,
            'write_timeout': self.write_timeout
        }
    
    def get_pool_params(self) -> Dict[str, Any]:
        """获取连接池参数"""
        params = self.get_connection_params()
        params.update({
            'maxconnections': self.max_connections,
            'mincached': self.min_connections,
            'maxcached': self.max_connections,
            'maxshared': 0,
            'blocking': True,
            'maxusage': None,
            'setsession': [],
            'ping': 0
        })
        return params


class MySQLConnectionManager:
    """MySQL连接管理器
    
    提供连接池管理、连接获取和释放功能
    """
    
    def __init__(self, config: DatabaseConfig):
        self.config = config
        self.logger = logging.getLogger(__name__)
        self._pool: Optional[pooling.PooledDB] = None
        self._initialize_pool()
    
    def _initialize_pool(self):
        """初始化连接池"""
        try:
            # 使用DBUtils的PooledDB创建连接池
            from dbutils.pooled_db import PooledDB
            
            self._pool = PooledDB(
                creator=pymysql,
                **self.config.get_pool_params()
            )
            
            self.logger.info(f"MySQL连接池初始化成功，最大连接数: {self.config.max_connections}")
            
        except ImportError:
            # 如果没有安装DBUtils，使用简单的连接管理
            self.logger.warning("DBUtils未安装，使用简单连接管理")
            self._pool = None
        except Exception as e:
            self.logger.error(f"连接池初始化失败: {e}")
            raise
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接"""
        connection = None
        try:
            if self._pool:
                connection = self._pool.connection()
            else:
                connection = pymysql.connect(**self.config.get_connection_params())
            
            yield connection
            
        except Exception as e:
            if connection:
                connection.rollback()
            self.logger.error(f"数据库连接错误: {e}")
            raise
        finally:
            if connection:
                connection.close()
    
    def test_connection(self) -> bool:
        """测试数据库连接"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                return result is not None
        except Exception as e:
            self.logger.error(f"数据库连接测试失败: {e}")
            return False
    
    def get_pool_status(self) -> Dict[str, Any]:
        """获取连接池状态"""
        if not self._pool:
            return {'pool_enabled': False}
        
        try:
            # 注意：DBUtils的PooledDB没有直接的状态查询方法
            # 这里返回配置信息
            return {
                'pool_enabled': True,
                'max_connections': self.config.max_connections,
                'min_connections': self.config.min_connections,
                'host': self.config.host,
                'database': self.config.database
            }
        except Exception as e:
            self.logger.error(f"获取连接池状态失败: {e}")
            return {'pool_enabled': False, 'error': str(e)}
    
    def close_pool(self):
        """关闭连接池"""
        if self._pool:
            try:
                self._pool.close()
                self.logger.info("连接池已关闭")
            except Exception as e:
                self.logger.error(f"关闭连接池失败: {e}")


class DatabaseInitializer:
    """数据库初始化器
    
    负责创建数据库、表结构和初始数据
    """
    
    def __init__(self, connection_manager: MySQLConnectionManager):
        self.connection_manager = connection_manager
        self.logger = logging.getLogger(__name__)
    
    def create_database_if_not_exists(self) -> bool:
        """创建数据库（如果不存在）"""
        try:
            # 临时连接到mysql数据库来创建目标数据库
            temp_config = DatabaseConfig(
                host=self.connection_manager.config.host,
                port=self.connection_manager.config.port,
                user=self.connection_manager.config.user,
                password=self.connection_manager.config.password,
                database='mysql',  # 连接到mysql系统数据库
                charset=self.connection_manager.config.charset
            )
            
            connection = pymysql.connect(**temp_config.get_connection_params())
            cursor = connection.cursor()
            
            # 创建数据库
            database_name = self.connection_manager.config.database
            cursor.execute(f"""
                CREATE DATABASE IF NOT EXISTS `{database_name}` 
                CHARACTER SET utf8mb4 
                COLLATE utf8mb4_unicode_ci
            """)
            
            connection.commit()
            connection.close()
            
            self.logger.info(f"数据库 '{database_name}' 创建成功或已存在")
            return True
            
        except Exception as e:
            self.logger.error(f"创建数据库失败: {e}")
            return False
    
    def create_tables(self) -> bool:
        """创建所有表结构"""
        try:
            with self.connection_manager.get_connection() as conn:
                cursor = conn.cursor()
                
                # 创建brain_wave表
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS brain_wave (
                        id VARCHAR(36) PRIMARY KEY,
                        session_id VARCHAR(100) NOT NULL,
                        room_id VARCHAR(50) NOT NULL,
                        user_id VARCHAR(50),
                        start_time DATETIME NOT NULL,
                        end_time DATETIME NOT NULL,
                        total_seconds INT NOT NULL,
                        metadata JSON,
                        raw_delta JSON,
                        raw_theta JSON,
                        raw_low_alpha JSON,
                        raw_high_alpha JSON,
                        raw_low_beta JSON,
                        raw_high_beta JSON,
                        raw_low_gamma JSON,
                        raw_mid_gamma JSON,
                        filtered_delta JSON,
                        filtered_theta JSON,
                        filtered_low_alpha JSON,
                        filtered_high_alpha JSON,
                        filtered_low_beta JSON,
                        filtered_high_beta JSON,
                        filtered_low_gamma JSON,
                        filtered_mid_gamma JSON,
                        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        INDEX idx_session_id (session_id),
                        INDEX idx_room_id (room_id),
                        INDEX idx_user_id (user_id),
                        INDEX idx_start_time (start_time)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                """)
                
                # 创建sleep_monitoring_metrics表
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS sleep_monitoring_metrics (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        session_id VARCHAR(100) NOT NULL,
                        room_id VARCHAR(50) NOT NULL,
                        start_time DATETIME NOT NULL,
                        total_seconds INT NOT NULL,
                        dominant_state VARCHAR(20),
                        avg_confidence DECIMAL(5,4),
                        sleep_stage_durations JSON,
                        data_delta_ratio JSON,
                        data_theta_ratio JSON,
                        data_low_alpha_ratio JSON,
                        data_high_alpha_ratio JSON,
                        data_low_beta_ratio JSON,
                        data_high_beta_ratio JSON,
                        data_low_gamma_ratio JSON,
                        data_mid_gamma_ratio JSON,
                        data_states JSON,
                        data_confidence JSON,
                        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                        INDEX idx_session_id (session_id),
                        INDEX idx_room_id (room_id),
                        INDEX idx_start_time (start_time)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                """)
                
                # 创建实时数据表
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS android_client_data (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        room_id VARCHAR(50) NOT NULL,
                        user_id VARCHAR(50),
                        session_id VARCHAR(100) NOT NULL,
                        timestamp DATETIME NOT NULL,
                        raw_delta FLOAT,
                        raw_theta FLOAT,
                        raw_low_alpha FLOAT,
                        raw_high_alpha FLOAT,
                        raw_low_beta FLOAT,
                        raw_high_beta FLOAT,
                        raw_low_gamma FLOAT,
                        raw_mid_gamma FLOAT,
                        filtered_delta FLOAT,
                        filtered_theta FLOAT,
                        filtered_low_alpha FLOAT,
                        filtered_high_alpha FLOAT,
                        filtered_low_beta FLOAT,
                        filtered_high_beta FLOAT,
                        filtered_low_gamma FLOAT,
                        filtered_mid_gamma FLOAT,
                        sleep_stage VARCHAR(20),
                        confidence DECIMAL(5,4),
                        device_info JSON,
                        signal_quality FLOAT,
                        received_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                        processed BOOLEAN DEFAULT FALSE,
                        INDEX idx_room_id (room_id),
                        INDEX idx_session_id (session_id),
                        INDEX idx_timestamp (timestamp),
                        INDEX idx_processed (processed)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                """)
                
                conn.commit()
                self.logger.info("所有表结构创建成功")
                return True
                
        except Exception as e:
            self.logger.error(f"创建表结构失败: {e}")
            return False
    
    def initialize_database(self) -> bool:
        """完整的数据库初始化"""
        try:
            # 1. 创建数据库
            if not self.create_database_if_not_exists():
                return False
            
            # 2. 创建表结构
            if not self.create_tables():
                return False
            
            self.logger.info("数据库初始化完成")
            return True
            
        except Exception as e:
            self.logger.error(f"数据库初始化失败: {e}")
            return False


# 全局连接管理器实例
_connection_manager: Optional[MySQLConnectionManager] = None


def get_connection_manager(config: Optional[DatabaseConfig] = None) -> MySQLConnectionManager:
    """获取全局连接管理器实例"""
    global _connection_manager
    
    if _connection_manager is None:
        if config is None:
            config = DatabaseConfig.from_env()
        _connection_manager = MySQLConnectionManager(config)
    
    return _connection_manager


def initialize_database(config: Optional[DatabaseConfig] = None) -> bool:
    """初始化数据库的便捷函数"""
    connection_manager = get_connection_manager(config)
    initializer = DatabaseInitializer(connection_manager)
    return initializer.initialize_database()


def test_database_connection(config: Optional[DatabaseConfig] = None) -> bool:
    """测试数据库连接的便捷函数"""
    connection_manager = get_connection_manager(config)
    return connection_manager.test_connection()


# 示例使用
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 从环境变量创建配置
    config = DatabaseConfig.from_env()
    print(f"数据库配置: {config.host}:{config.port}/{config.database}")
    
    # 测试连接
    if test_database_connection(config):
        print("数据库连接测试成功")
    else:
        print("数据库连接测试失败")
    
    # 初始化数据库
    if initialize_database(config):
        print("数据库初始化成功")
    else:
        print("数据库初始化失败")